# -*- coding: utf-8 -*-
# @Time     : 2020/7/7 2:28 下午
# @Author   : huangxiong
# @FileName : risk_decomposition.py
# @Comment  : 风险分解的核心函数
# @Software : PyCharm

import pandas as pd
import numpy as np

from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.project_tool.time_tool import date_shifter
from quant_researcher.quant.project_tool.math_func.index_fitting import process
from quant_researcher.quant.datasource_fetch.factor_api import factor_return_related
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import \
    BARRA_FACTOR_NAME_DICT, BOND_FACTOR_NAME_DICT, BARRA_FULL_FACTOR
from quant_researcher.quant.project_tool.exception import FOFUserError, ERR_U_1

INDEX_COMPONENTS_TABLE = 'idx_components_weight'
STOCK_BARRA_EXPOSURE_TABLE = 'stk_barra_exposure'
STOCK_BARRA_COVARIANCE_TABLE = 'stk_barra_covariance_day'
INDEX_BASEINFO_TABLE = 'idx_baseinfo'
FUND_STOCK_PORTFOLIO_TABLE = 'mf_iv_stockportfolio'

FACTOR_NAME_DICT = BARRA_FACTOR_NAME_DICT.copy()
FACTOR_NAME_DICT.update(BOND_FACTOR_NAME_DICT)

BARRA_FACTOR_TYPE = 'BARRA_FYZ'
BOND_FACTOR_TYPE = 'BOND_7YZ'
# 这里的Default没有``符号
BOND_7_FACTOR = ['Level', 'Slope', 'Convertible', 'Credit', 'Convex', 'Default', 'Currency']


def risk_decomposition(calc_date, fund_list, weight_list, customer_create, index_code, customer_index, calc_model):
    """

    :param str calc_date: 计算日期，格式形如'2020-07-09'
    :param str fund_list: 基金列表，如 ['110022', '001410', '003095', '008817']
    :param str weight_list: 基金的权重列表，如 [0.2, 0.3, 0.4, 0.1]
    :param customer_create: 是否是用户自建指数，1-是，0-否，默认0
    :param index_code: str，基准指数代码
    :param customer_index: 用户自建指数，customer_create=1时必须传入该参数。
                           需要传入构建自建指数所需要的指数代码及权重，格式{"CBA00301.CS": 0.5, "000300": 0.5}
    :param calc_model: str，计算模式，可选股票因子('s')、债券因子('b')、股债混合因子('s&b')
    :return: 组合和基准的每个因子的风险暴露、边际风险贡献、总风险贡献
    """
    if (customer_create == 1 and customer_index is None) or \
            (customer_create == 0 and index_code is None):
        raise FOFUserError(ERR_U_1, "如果customer_create为1，需要输入customer_index，"
                                    "如果customer_create为0，需要输入index_code，"
                                    "请按要求重新输入")

    calc_date_t = calc_date.replace('-', '')

    conn_base = db_conn.get_basic_data_conn()
    conn_factor = db_conn.get_tk_factors_conn()

    fund_weight_df = pd.DataFrame({'fund_code': fund_list, 'weight': weight_list})

    # 将基金和指数代码组装，作为传入process函数的入参
    fund_bench_df = pd.DataFrame({'fund_code': fund_list})
    fund_bench_df['index_code'] = index_code
    fund_bench_list = list(fund_bench_df.values)

    s_res_df = b_res_df = None
    # 计算barra因子暴露
    if calc_model in ['s', 's&b']:
        # 股票的barra因子暴露
        stock_exposure_df = pd.read_sql(f"select * "
                                        f"from {STOCK_BARRA_EXPOSURE_TABLE} "
                                        f"where tradedate = '{calc_date}' ", conn_factor)
        stock_exposure_df = stock_exposure_df.drop(columns=['tid', 'tradedate', 'sw_ind1',
                                                            'titime', 'tutime', 'tutimestamp'])
        stock_exposure_df = stock_exposure_df.rename(columns={'stockcode': 'stock_code'})
        # 股票的barra因子收益率协方差矩阵
        stock_factor_cov_df = pd.read_sql(f"select * "
                                          f"from {STOCK_BARRA_COVARIANCE_TABLE} "
                                          f"where tradedate = '{calc_date_t}' ", conn_factor)
        stock_factor_cov_df = stock_factor_cov_df.drop(columns=['tid', 'tradedate', 'titime', 'tutime', 'tutimestamp'])
        stock_factor_cov_df = stock_factor_cov_df.set_index('factor')
        # 数据库中协方差值乘以了1e6，这里调回去
        stock_factor_cov_df = stock_factor_cov_df / 1e6
        stock_factor_cov_df = stock_factor_cov_df.reindex(BARRA_FULL_FACTOR)[BARRA_FULL_FACTOR]
        stock_factor_cov_df = stock_factor_cov_df.fillna(0)
        stock_factor_cov_df = stock_factor_cov_df.values

        fund_where_condition = get_lst_condition(fund_list, 'fund_code')
        # 去找基金最近的股票持仓数据
        fund_hold_df = pd.read_sql(f"select a.fund_code, a.stock_code, a.hold_mktvalue as value, "
                                   f"a.pptinnv as stock_weight "
                                   f"from {FUND_STOCK_PORTFOLIO_TABLE} as a, "
                                   f"(select fund_code, max(rept_enddate) as rept_enddate "
                                   f"from {FUND_STOCK_PORTFOLIO_TABLE} "
                                   f"where rept_enddate <= '{calc_date_t}' "
                                   f"and rept_type in (5, 6) "
                                   f"and {fund_where_condition} "
                                   f"group by fund_code) as b "
                                   f"where a.fund_code = b.fund_code "
                                   f"and a.rept_enddate = b.rept_enddate "
                                   f"and a.rept_type in (5, 6) ", conn_base)
        # 计算组合的股票持仓比例
        p_stock_weight = fund_hold_df.groupby('fund_code')['stock_weight'].sum().reset_index()
        p_stock_weight = p_stock_weight.merge(fund_weight_df, how='left', on='fund_code')
        p_stock_weight = (p_stock_weight['stock_weight'] * p_stock_weight['weight']).sum()
        # 如果基金有股票持仓数据，且股票持仓比例不小于1%，则用股票持仓数据计算因子暴露
        if not fund_to_remove(fund_hold_df, 'value') and p_stock_weight >= 1:
            f_stock_df = fund_hold_df.groupby('stock_code')['value'].sum().reset_index()
            f_stock_df['weight'] = f_stock_df['value'] / f_stock_df['value'].sum()
            f_stock_df = f_stock_df.drop(columns='value')
            f_stock_df = f_stock_df.merge(stock_exposure_df, how='left', on='stock_code')
            f_stock_df = f_stock_df.fillna(0)
        # 否则，用回归计算因子暴露
        else:
            # 用process函数计算基金的barra因子暴露
            f_stock_df = process(code_class='fund', today=calc_date, daily_or_monthly='D',
                                 i_or_s='barra', i_or_s_specific=BARRA_FACTOR_TYPE,
                                 test_codes=fund_bench_list, data_4_api=True, process_number=5)
            f_stock_df = get_factor_exposure('fund_code', fund_list, 'barra', f_stock_df)

            f_stock_df = fund_weight_df.merge(f_stock_df, how='left', on='fund_code')
            f_stock_df = f_stock_df.fillna(0)
            f_stock_df = f_stock_df[['fund_code', 'weight'] + BARRA_FULL_FACTOR]
        p_s_exposure, p_s_mrc, p_s_trc = calc_risk(f_stock_df, stock_factor_cov_df, calc_model='s')
        if p_s_mrc is None:
            return

        if customer_create == 0:
            i_s_exposure, i_s_mrc, i_s_trc = get_index_stock_risk_res(index_code, calc_date, conn_base,
                                                                      stock_exposure_df, stock_factor_cov_df)
            if i_s_mrc is None:
                return
        else:
            i_s_exposure_list = []
            i_s_mrc_list = []
            i_s_trc_list = []

            for i in customer_index.keys():
                i_s_exposure_tmp, i_s_mrc_tmp, i_s_trc_tmp = \
                    get_index_stock_risk_res(i, calc_date, conn_base,
                                             stock_exposure_df, stock_factor_cov_df)
                if i_s_mrc_tmp is not None:
                    i_s_exposure_tmp = i_s_exposure_tmp * customer_index[i]
                    i_s_exposure_list.append(i_s_exposure_tmp)
                    i_s_mrc_tmp = i_s_mrc_tmp * customer_index[i]
                    i_s_mrc_list.append(i_s_mrc_tmp)
                    i_s_trc_tmp = i_s_trc_tmp * customer_index[i]
                    i_s_trc_list.append(i_s_trc_tmp)

            if len(i_s_mrc_list) == 0:
                return

            i_s_exposure = np.sum(i_s_exposure_list, axis=0)
            i_s_mrc = np.sum(i_s_mrc_list, axis=0)
            i_s_trc = np.sum(i_s_trc_list, axis=0)

        s_res_df = get_res_df(p_s_exposure, p_s_mrc, p_s_trc, i_s_exposure, i_s_mrc, i_s_trc, BARRA_FULL_FACTOR)

    if calc_model in ['b', 's&b']:
        # 计算债券因子收益率的协方差矩阵，用近一年的收益率序列
        calc_date_before = date_shifter(before=calc_date, step='years', how_many=-1)

        bond_factor_ret = factor_return_related.get_bond_factor_return(
            BOND_FACTOR_TYPE, calc_date_before, calc_date
        )
        bond_factor_ret = bond_factor_ret.drop(columns='end_date')
        bond_factor_cov_df = bond_factor_ret.cov()
        bond_factor_cov_df = bond_factor_cov_df.reindex(BOND_7_FACTOR)[BOND_7_FACTOR]
        bond_factor_cov_df = bond_factor_cov_df.fillna(0)
        bond_factor_cov_df = bond_factor_cov_df.values

        # 用process函数回归计算基金和指数的bond因子暴露
        f_b_df = process(code_class='fund', today=calc_date, daily_or_monthly='D',
                         i_or_s='bond_factors', i_or_s_specific=BOND_FACTOR_TYPE,
                         test_codes=fund_bench_list, data_4_api=True, process_number=5)
        f_b_df = get_factor_exposure('fund_code', fund_list, 'bond', f_b_df)
        f_b_df = fund_weight_df.merge(f_b_df, how='left', on='fund_code')
        f_b_df = f_b_df.fillna(0)
        f_b_df = f_b_df[['fund_code', 'weight'] + BOND_7_FACTOR]

        if customer_create == 0:
            i_b_df = get_index_bond_risk_res(index_code, calc_date)
        else:
            i_b_df_list = []
            for i in customer_index.keys():
                i_b_df_tmp = get_index_bond_risk_res(i, calc_date)
                i_b_df_list.append(i_b_df_tmp)
            i_b_df = pd.concat(i_b_df_list)
            i_b_df = i_b_df.drop(columns=['index_code'])
            i_b_df = i_b_df.mul(np.array(list(customer_index.values())), axis=0)
            i_b_df = i_b_df.sum().to_frame().T

        p_b_exposure, p_b_mrc, p_b_trc = calc_risk(f_b_df, bond_factor_cov_df, calc_model='b')
        if p_b_mrc is None:
            return
        i_b_exposure, i_b_mrc, i_b_trc = calc_risk(i_b_df, bond_factor_cov_df, calc_model='b')
        if i_b_mrc is None:
            return

        b_res_df = get_res_df(p_b_exposure, p_b_mrc, p_b_trc, i_b_exposure, i_b_mrc, i_b_trc, BOND_7_FACTOR)

    if calc_model == 's':
        res_df = s_res_df
    elif calc_model == 'b':
        res_df = b_res_df
    elif calc_model == 's&b':
        res_df = pd.concat([s_res_df, b_res_df], ignore_index=True, sort=False)
    else:
        raise NotImplementedError

    factor_name_df = pd.DataFrame(list(FACTOR_NAME_DICT.items()), columns=['factor_code', 'factor_name'])
    res_df = res_df.merge(factor_name_df, how='left', on='factor_code')

    conn_base.close()
    conn_factor.close()

    return res_df


def get_factor_exposure(code_type, code_list, factor_type, data_df):
    if factor_type == 'barra':
        factor_list = BARRA_FULL_FACTOR
    else:
        factor_list = BOND_7_FACTOR

    if data_df.empty:
        exposure_df = pd.DataFrame({code_type: code_list})
        for factor_name in factor_list:
            exposure_df[factor_name] = 0
    else:
        exposure_df = data_df[[code_type, 'factor_code', 'factor_beta']]
        exposure_df = exposure_df.set_index(
            [code_type, 'factor_code'])['factor_beta'].unstack().reset_index()

    return exposure_df


def calc_risk(calc_df, factor_cov_array, calc_model):
    weight_array = calc_df['weight'].values
    weight_array = weight_array.reshape(-1, 1)
    if calc_model == 's':
        factor_exposure_array = calc_df[BARRA_FULL_FACTOR].values.T
    else:
        factor_exposure_array = calc_df[BOND_7_FACTOR].values.T
    factor_weight = np.dot(factor_exposure_array, weight_array)

    mrc, trc = risk_contribution(factor_weight, factor_cov_array)

    return factor_weight, mrc, trc


def risk_contribution(weight_array, cov_array):
    """

    :param weight_array: 资产权重序列
    :param cov_array: 资产收益协方差矩阵
    :return: 资产边际风险贡献，总风险贡献
    """

    # 组合风险
    sigma_p = np.dot(weight_array.T, cov_array)
    sigma_p = np.dot(sigma_p, weight_array)
    sigma_p = np.sqrt(sigma_p[0, 0])

    if sigma_p == 0:
        LOG.error('组合风险为0，应该是哪里出了问题')
        mrc = trc = None
    else:
        # 边际风险贡献
        mrc = np.dot(cov_array, weight_array) / sigma_p
        # 总风险贡献
        trc = mrc * weight_array
    return mrc, trc


def get_res_df(p_exposure, p_mrc, p_trc, i_exposure, i_mrc, i_trc, factor_list):
    # 输出组合和基准的每个因子的风险暴露、边际风险贡献、总风险贡献
    # 以及主动风险暴露、主动边际贡献、主动总风险贡献
    res_df = pd.DataFrame({'factor_code': factor_list,
                           'portfolio_exposure': p_exposure[:, 0],
                           'index_exposure': i_exposure[:, 0],
                           'portfolio_mrc': p_mrc[:, 0], 'index_mrc': i_mrc[:, 0],
                           'portfolio_trc': p_trc[:, 0], 'index_trc': i_trc[:, 0]})
    res_df = res_df.set_index('factor_code')
    res_df['active_exposure'] = res_df['portfolio_exposure'] - res_df['index_exposure']
    res_df['active_mrc'] = res_df['portfolio_mrc'] - res_df['index_mrc']
    res_df['active_trc'] = res_df['portfolio_trc'] - res_df['index_trc']
    res_df = res_df.round(4)
    res_df = res_df.reset_index()
    return res_df


def fund_to_remove(df, col_name):
    if df.empty or df.dropna().empty or df[df[col_name] > 0].empty:
        return True
    return False


def get_index_stock_risk_res(index_code, calc_date, conn_base, stock_exposure_df, stock_factor_cov_df):
    calc_date_t = calc_date.replace('-', '')
    # 找到指数在离end_date这天最近的成份股及权重
    index_latest_weight_date = pd.read_sql(f"select max(trade_date) "
                                           f"from {INDEX_COMPONENTS_TABLE} "
                                           f"where index_code = '{index_code}' "
                                           f"and trade_date <= '{calc_date_t}' ", conn_base).iloc[0, 0]

    # 找到所有债券指数的列表
    bond_index_list = pd.read_sql(f"select index_code "
                                  f"from {INDEX_BASEINFO_TABLE} "
                                  f"where index_type = '03' ", conn_base)['index_code'].tolist()

    # 如果能找到指数的成份股，且指数不是债券指数，则用成份股计算
    if (index_latest_weight_date is not None) and (index_code not in bond_index_list):
        index_weight_df = pd.read_sql(f"select secu_code as stock_code, weight/100 as weight "
                                      f"from {INDEX_COMPONENTS_TABLE} "
                                      f"where index_code = '{index_code}' "
                                      f"and trade_date = '{index_latest_weight_date}' ", conn_base)
        # # 如果指数权重最新日期比现在日期小，按照股票涨跌幅调整个股权重
        # if index_latest_weight_date < calc_date_t:
        #     stock_ret = stock_price_related.get_stock_total_return(index_latest_weight_date, calc_date_t)
        #     index_weight_df = index_weight_df.merge(stock_ret, how='left', on='stock_code').fillna(0)
        #     index_weight_df['weight2'] = index_weight_df['weight'] * (1 + index_weight_df['ret'])
        #     index_weight_df['weight_adj'] = index_weight_df['weight2'] / index_weight_df['weight2'].sum()
        #     index_weight_df = index_weight_df[
        #         ['stock_code', 'weight_adj']].rename(columns={'weight_adj': 'weight'})

        i_stock_df = index_weight_df.merge(stock_exposure_df, how='left', on='stock_code')
        if i_stock_df.dropna().empty:
            LOG.error(f"{index_code} 该指数的成份股没有股票，可能是债券指数，且在指数信息表中没有分类")
            return None, None, None
        i_stock_df = i_stock_df.fillna(0)
    # 如果找不到指数的成份股，或者指数为债券指数，则去计算指数的因子暴露
    else:
        index_weight_df = pd.DataFrame({'index_code': [index_code], 'weight': 1})
        i_stock_df = process(code_class='index', today=calc_date, daily_or_monthly='D',
                             i_or_s='barra', i_or_s_specific=BARRA_FACTOR_TYPE,
                             test_codes=[index_code], data_4_api=True, process_number=1)
        if i_stock_df.empty:
            i_stock_df = pd.DataFrame(0, index=[index_code], columns=BARRA_FULL_FACTOR)
            i_stock_df = i_stock_df.reset_index().rename(columns={'index': 'index_code'})
        else:
            i_stock_df = get_factor_exposure('index_code', [index_code], 'barra', i_stock_df)

        i_stock_df = index_weight_df.merge(i_stock_df, how='left', on='index_code')
        i_stock_df = i_stock_df.fillna(0)
        i_stock_df = i_stock_df[['index_code', 'weight'] + BARRA_FULL_FACTOR]
    i_s_exposure, i_s_mrc, i_s_trc = calc_risk(i_stock_df, stock_factor_cov_df, calc_model='s')

    return i_s_exposure, i_s_mrc, i_s_trc


def get_index_bond_risk_res(index_code, calc_date):
    i_b_df = process(code_class='index', today=calc_date, daily_or_monthly='D',
                     i_or_s='bond_factors', i_or_s_specific=BOND_FACTOR_TYPE,
                     test_codes=[index_code], data_4_api=True, process_number=1)
    index_weight_df = pd.DataFrame({'index_code': [index_code], 'weight': 1})
    if i_b_df.empty:
        i_b_df = pd.DataFrame(0, index=[index_code], columns=BOND_7_FACTOR)
        i_b_df = i_b_df.reset_index().rename(columns={'index': 'index_code'})
    else:
        i_b_df = get_factor_exposure('index_code', [index_code], 'bond', i_b_df)
        i_b_df = index_weight_df.merge(i_b_df, how='left', on='index_code')
        i_b_df = i_b_df.fillna(0)
        i_b_df = i_b_df[['index_code', 'weight'] + BOND_7_FACTOR]

    return i_b_df


if __name__ == "__main__":
    risk_decomposition(calc_date='2020-08-04', fund_list=['110022', '001410', '003095', '008817'],
                       weight_list=[0.2, 0.3, 0.4, 0.1], index_code='000300', calc_model='s&b')
